import unittest
from warp_field import profile
import time

import torch.nn as nn

class ModuleTest(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x):
        print("forward ModuleTest")
        time.sleep(0.5)
    
    def backward(self, x):
        print("backward ModuleTest")
        time.sleep(0.5)

class NVTXTest(unittest.TestCase):
    def setUp(self) -> None:
        print("########## setup NVTXTest ##########")

    def tearDown(self) -> None:
        print("########## teardown NVTXTest ##########")

    def test_nvtx(self):
        module = ModuleTest()
        wrapper = profile.NVTXModuleWrapper(module)
        wrapper.register()
        module(1)
        wrapper._backward_pre_hook_fn(module, None, None)
        module.backward(1)
        wrapper._backward_post_hook_fn(module, None, None)